/*
 * Copyright (c) 2023, Arm Limited and Contributors. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

extern "C" {
#include "aws_network_manager/aws_network_manager.h"
#include "backoff_algorithm.h"
#include "iot_socket.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ssl.h"
}

#include "fff.h"
#include "gtest/gtest.h"

DEFINE_FFF_GLOBALS

class TestAwsNetworkManager : public ::testing::Test {
public:
    TestAwsNetworkManager()
    {
        RESET_FAKE(mbedtls_entropy_init);
        RESET_FAKE(mbedtls_entropy_free);
        RESET_FAKE(mbedtls_ctr_drbg_seed);
        RESET_FAKE(mbedtls_ssl_handshake);
        RESET_FAKE(mbedtls_ssl_write);
        RESET_FAKE(mbedtls_ssl_read);
        RESET_FAKE(mbedtls_ssl_close_notify);
        RESET_FAKE(iotSocketCreate);
        RESET_FAKE(iotSocketClose);
        RESET_FAKE(iotSocketGetHostByName);
        RESET_FAKE(BackoffAlgorithm_GetNextBackoff);
    }
};

TEST_F(TestAwsNetworkManager, network_initialization_fails_when_context_is_nullptr)
{
    AwsNetworkContextConfig_t network_config = {};
    AwsNetwork_init(nullptr, &network_config);

    // Connection should fail if initialization was not successful
    EXPECT_NE(0, AwsNetwork_connect(nullptr, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_initialization_fails_when_config_is_nullptr)
{
    NetworkContext_t network_context = {};
    AwsNetwork_init(&network_context, nullptr);

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_initialization_succeeds_when_config_parameters_are_provided)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_EQ(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_when_ctx_is_nullptr)
{
    EXPECT_EQ(-1, AwsNetwork_connect(nullptr, [](NetworkContext_t *context) {
        (void)context;
        return 0;
    }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_if_configuring_tls_contexts_fails)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    AwsNetwork_init(&network_context, &network_config);

    // Prevents successful TLS context configuration
    mbedtls_ctr_drbg_seed_fake.return_val = -1;

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_when_configuring_socket_fails)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    // Prevents successful socket configuration
    iotSocketCreate_fake.return_val = -1;

    BackoffAlgorithm_GetNextBackoff_fake.return_val = BackoffAlgorithmRetriesExhausted;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_when_server_connection_fails)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    // Prevents successful server connection
    iotSocketGetHostByName_fake.return_val = -1;

    BackoffAlgorithm_GetNextBackoff_fake.return_val = BackoffAlgorithmRetriesExhausted;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_when_tls_connection_fails)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    // Prevents successful tls connection
    mbedtls_ssl_handshake_fake.return_val = -1;

    BackoffAlgorithm_GetNextBackoff_fake.return_val = BackoffAlgorithmRetriesExhausted;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_EQ(-1, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
        (void)context;
        return 0;
    }));
}

TEST_F(TestAwsNetworkManager, network_resets_when_ssl_session_cannot_be_reset)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    mbedtls_ssl_handshake_fake.return_val = -1;
    mbedtls_ssl_session_reset_fake.return_val = -1;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_fails_when_callback_fails)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    BackoffAlgorithm_GetNextBackoff_fake.return_val = BackoffAlgorithmRetriesExhausted;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return -1;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_succeeds_when_server_and_tls_connection_is_successful)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_EQ(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, network_connection_succeeds_even_when_callback_is_not_provided)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_EQ(0, AwsNetwork_connect(&network_context, nullptr));
}

TEST_F(TestAwsNetworkManager, sending_message_fails_when_ctx_is_nullptr)
{
    uint8_t buffer = 0;

    EXPECT_NE(0, AwsNetwork_send(nullptr, &buffer, sizeof(buffer)));
}

TEST_F(TestAwsNetworkManager, sending_message_fails_when_buf_is_nullptr)
{
    NetworkContext_t network_context = {};

    EXPECT_NE(0, AwsNetwork_send(&network_context, nullptr, 0));
}

TEST_F(TestAwsNetworkManager, sending_message_succeeds_when_message_is_sent)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    AwsNetwork_init(&network_context, &network_config);

    uint8_t buffer = 0;

    mbedtls_ssl_write_fake.return_val = MBEDTLS_ERR_SSL_WANT_WRITE;

    EXPECT_EQ(0, AwsNetwork_send(&network_context, &buffer, sizeof(buffer)));
}

TEST_F(TestAwsNetworkManager, receiving_message_fails_when_ctx_is_nullptr)
{
    uint8_t buffer = 0;

    EXPECT_NE(0, AwsNetwork_recv(nullptr, &buffer, sizeof(buffer)));
}

TEST_F(TestAwsNetworkManager, receiving_message_fails_when_buf_is_nullptr)
{
    NetworkContext_t network_context = {};

    EXPECT_NE(0, AwsNetwork_recv(&network_context, nullptr, 0));
}

TEST_F(TestAwsNetworkManager, receiving_message_succeeds_when_message_is_received)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    AwsNetwork_init(&network_context, &network_config);

    uint8_t buffer = 0;

    mbedtls_ssl_read_fake.return_val = MBEDTLS_ERR_SSL_WANT_READ;

    EXPECT_EQ(0, AwsNetwork_recv(&network_context, &buffer, sizeof(buffer)));
}

TEST_F(TestAwsNetworkManager, closing_network_fails_when_context_is_nullptr)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    network_context.socket_fd = IOT_SOCKET_EISCONN;
    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    AwsNetwork_close(nullptr);

    // If the network has not been closed, it should still be possible to connect to the network.
    EXPECT_EQ(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}

TEST_F(TestAwsNetworkManager, closing_network_succeeds_when_config_parameters_are_provided)
{
    NetworkContext_t network_context = {};
    AwsNetworkContextConfig_t network_config = {};
    network_config.server_certificate = "";
    network_config.client_certificate = "";
    network_config.client_private_key = "";
    network_config.server_name = "";
    AwsNetwork_init(&network_context, &network_config);

    mbedtls_ctr_drbg_seed_fake.return_val = 0;
    mbedtls_x509_crt_parse_fake.return_val = 0;
    mbedtls_ssl_config_defaults_fake.return_val = 0;
    mbedtls_pk_parse_key_fake.return_val = 0;
    mbedtls_ssl_conf_own_cert_fake.return_val = 0;
    mbedtls_ssl_setup_fake.return_val = 0;
    mbedtls_ssl_set_hostname_fake.return_val = 0;

    EXPECT_EQ(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));

    AwsNetwork_close(&network_context);

    // It should not be possible to connect without initialising the network again
    EXPECT_NE(0, AwsNetwork_connect(&network_context, [](NetworkContext_t *context) {
                  (void)context;
                  return 0;
              }));
}
